﻿using System.Collections;
using System.Collections.Generic;
using UnityEngine;


namespace FFTOcean
{
    ///<summary>
    /// 使用Stockham FFT算法
    /// 代码从 http://wwwa.pikara.ne.jp/okojisan/otfft-en/stockham2.html list7的c++翻译过来
    ///</summary>
    public class FFT
    {
        ///<summary>
        /// y  : work area(or output sequence if eo == 1) 
        /// x  : input sequence(or output sequence if eo == 0) 
        /// eo : x is output if eo == 0, y is output if eo == 1 
        /// s  : stride 
        /// n  : sequence length
        ///</summary>
        static void fft0(int n, int s, bool eo, Complex[] x, Complex[] y)
        {

            int m = n / 2;
            float theta0 = 2 * Mathf.PI / n;
            if (n == 1)
            {
                if (eo) for (int q = 0; q < s; q++) y[q] = x[q];
            }
            else
            {
                for (int p = 0; p < m; p++)
                {
                    Complex wp = new Complex(Mathf.Cos(p * theta0), -Mathf.Sin(p * theta0));
                    for (int q = 0; q < s; q++)
                    {
                        Complex a = x[q + s * (p + 0)];
                        Complex b = x[q + s * (p + m)];
                        y[q + s * (2 * p + 0)] = a + b;
                        y[q + s * (2 * p + 1)] = (a - b) * wp;
                    }
                }
                fft0(n / 2, 2 * s, !eo, y, x);
            }
        }

        ///<summary>
        /// x : input/output sequence 
        /// n : sequence length 
        /// Inverse Fourier transform
        ///</summary>
        public static void ifft(int n, Complex[] x)
        {
            for (int p = 0; p < n; p++) x[p] = x[p].Conjugate();
            Complex[] y = new Complex[n];
            fft0(n, 1, false, x, y);
            for (int k = 0; k < n; k++) x[k] = x[k].Conjugate();
        }
    }

    ///<summary>
    /// 负责频谱预计算和所有FFT计算
    /// 符号和原理借鉴  https://zhuanlan.zhihu.com/p/64414956
    ///</summary>
    public class OceanCalculation
    {
        private int N;
        private int seed = 0;
        private float A;
        private Vector2 OMEGA;
        private float V;
        private float g = 9.81f;
        private int Nplus1;

        ///<summary>
        /// 存储Phillips频谱表达式中可预计算的部分,形状[,,3], [,,0]存储h0，1存储h0*，2存储w(k)
        ///</summary>
        private Complex[,,] m_phillips;

        ///<summary>
        /// N: 海面patch大小，同时也是横向或纵向顶点个数
        /// A： 影响波浪振幅大小
        /// OMEGA： 风向
        /// V： 风速
        ///</summary>
        public OceanCalculation(int N, float A, Vector2 OMEGA, float V)
        {
            this.N = N;
            Nplus1 = N + 1;
            this.A = A;
            this.OMEGA = OMEGA;
            this.V = V;
            m_phillips = new Complex[N, N, 3];
            InitPhillips();
        }

        ///<summary>
        /// 预计算Phillips频谱
        ///</summary>
        private void InitPhillips()
        {
            Random.InitState(seed);
            float L = V * V / g;
            for (int i = 0; i < N; i++)
            {
                for (int j = 0; j < N; j++)
                {
                    float kx = 2 * Mathf.PI * ((-1 * N / 2) + i) / N;
                    float kz = 2 * Mathf.PI * ((-1 * N / 2) + j) / N;


                    float u1 = Random.Range(0f, 1f);
                    float u2 = Random.Range(0f, 1f);

                    //epsilon为服从正态分布的独立随机变量
                    float epsilon_r = Mathf.Sqrt(-2 * Mathf.Log(u1)) * Mathf.Cos(2 * Mathf.PI * u2);
                    float epsilon_i = Mathf.Sqrt(-2 * Mathf.Log(u1)) * Mathf.Sin(2 * Mathf.PI * u2);

                    float k = Mathf.Sqrt(kx * kx + kz * kz);

                    float P_h;
                    if (kx == 0 && kz == 0)
                    {
                        P_h = 0;
                    }
                    else
                    {
                        Vector2 omega = OMEGA.normalized;
                        Vector2 k_vec = new Vector2(kx, kz);
                        float k_length = k_vec.magnitude;

                        float k_dot_w = Vector2.Dot(omega, k_vec);
                        float k_dot_w2 = k_dot_w * k_dot_w;

                        float k2 = k * k;
                        float L2 = L * L;

                        float damping = 0.001f;
                        float l2 = L2 * damping * damping;

                        P_h = A * Mathf.Exp(-1 / (k2 * L2)) / Mathf.Pow(k, 4) * k_dot_w2 * Mathf.Exp(-k2 * l2);
                    }


                    Complex r = new Complex(epsilon_r, epsilon_i);
                    Complex h0 = 1 / Mathf.Sqrt(2) * r * Mathf.Sqrt(P_h);
                    Complex h0_conj = h0.Conjugate();

                    float w_0 = 2.0f * Mathf.PI / 200.0f;

                    m_phillips[i, j, 0] = h0;
                    m_phillips[i, j, 1] = h0_conj;
                    m_phillips[i, j, 2] = Mathf.Floor(Mathf.Sqrt(g * k) / w_0) * w_0;
                }
            }
        }

        ///<summary>
        /// 计算需要FFT的部分
        /// height： 高度
        /// displacement: Complex[2][N,N],x和z轴的位移，[0][i,j]为x轴位移，[1][i,j]为z轴
        /// normal： Complex[2][N,N],法线对x和z的梯度，[0][i,j]为x，[1][,]为z
        /// foam: Complex[3][N,N],白沫所需要的三个Jacoobi值中的偏导值，0对应Jxx，1 Jxz，2 Jzz
        /// t： 时间
        ///</summary>
        public void FFTCalculation(Complex[,] height, Complex[][,] displacement, Complex[][,] normal, Complex[][,] foam, float t)
        {
            getAllSpectrum(height, displacement, normal, foam, t);
            int num = 2;
            inplace2DFFT(height);
            for (int i = 0; i < num; i++)
            {
                inplace2DFFT(displacement[i]);
                inplace2DFFT(normal[i]);
            }
            for (int i = 0; i < 3; i++)
            {
                inplace2DFFT(foam[i]);
            }
        }


        ///<summary>
        /// 计算FFT的系数，对高度而言是频谱
        ///</summary>
        private void getAllSpectrum(Complex[,] height, Complex[][,] displacement, Complex[][,] normal, Complex[][,] foam, float t)
        {
            for (int i = 0; i < N; i++)
            {
                for (int j = 0; j < N; j++)
                {
                    //height
                    float kx = 2 * Mathf.PI * ((-1 * N / 2) + i) / N;
                    float kz = 2 * Mathf.PI * ((-1 * N / 2) + j) / N;
                    float omega_k_t = (float)m_phillips[i, j, 2] * t;
                    Complex exp_term1 = new Complex(Mathf.Cos(omega_k_t), Mathf.Sin(omega_k_t));
                    Complex exp_term2 = new Complex(Mathf.Cos(-omega_k_t), Mathf.Sin(-omega_k_t));
                    Complex term1 = m_phillips[i, j, 0] * exp_term1;
                    Complex term2 = m_phillips[i, j, 1] * exp_term2;
                    height[i, j] = term1 + term2;

                    //normal
                    Complex kx_i = new Complex(0, kx);
                    Complex kz_i = new Complex(0, kz);
                    normal[0][i, j] = term1 + term2 * kx_i;
                    normal[1][i, j] = term1 + term2 * kz_i;

                    //displacement
                    float k = Mathf.Sqrt(kx * kx + kz * kz);
                    if (k < 0.00001f)
                    {
                        displacement[0][i, j] = 0;
                        displacement[1][i, j] = 0;
                    }
                    else
                    {
                        displacement[0][i, j] = -1 / k * normal[0][i, j];
                        displacement[1][i, j] = -1 / k * normal[1][i, j];
                    }

                    //foam
                    foam[0][i, j] = kx_i * displacement[0][i, j];
                    foam[1][i, j] = kz_i * displacement[0][i, j];
                    foam[2][i, j] = kz_i * displacement[1][i, j];
                }
            }
        }

        ///<summary>
        ///就地二维IDFT
        ///result： 频谱
        ///</summary>
        public void inplace2DFFT(Complex[,] result)
        {
            for (int i = 1; i < N; i += 2)
            {
                for (int j = 0; j < N; j++)
                {
                    result[i, j] = result[i, j] * (-1);
                }
            }

            Complex[] tmp = new Complex[N];
            //IFFT,求C
            for (int i = 0; i < N; i++)
            {
                copyRowToArray(i, result, N, tmp);
                FFT.ifft(N, tmp);
                copyArrayToRow(i, result, N, tmp);
            }

            //求B，符号矫正2
            for (int i = 0; i < N; i++)
            {
                for (int j = 0; j < N; j++)
                {
                    result[i, j] = result[i, j] * Mathf.Pow(-1, i + j); 
                }
            }

            //IFFT,求A
            for (int i = 0; i < N; i++)
            {
                copyColumnToArray(i, result, N, tmp);
                FFT.ifft(N, tmp);
                copyArrayToColumn(i, result, N, tmp);
            }

            //计算高度h，符号矫正3
            for (int i = 0; i < N; i++)
            {
                for (int j = 1; j < N; j += 2)
                {
                    result[i, j] = result[i, j] * (-1);
                }
            }

            // 将边界填充为正确的值来支持tiling
            result[N, N] = result[0, 0];
            for (int i = 0; i < N; i++)
            {
                result[i, N] = result[i, 0];
            }
            for (int y = 0; y < N; y++)
            {
                result[N, y] = result[0, y];
            }

        }

        private void copyRowToArray(int index, Complex[,] src, int length, Complex[] dest)
        {
            for (int i = 0; i < length; i++)
            {
                dest[i] = src[index, i];
            }
        }

        private void copyArrayToRow(int index, Complex[,] dest, int length, Complex[] src)
        {
            for (int i = 0; i < length; i++)
            {
                dest[index, i] = src[i];
            }
        }

        private void copyColumnToArray(int index, Complex[,] src, int length, Complex[] dest)
        {
            for (int i = 0; i < length; i++)
            {
                dest[i] = src[i, index];
            }
        }

        private void copyArrayToColumn(int index, Complex[,] dest, int length, Complex[] src)
        {
            for (int i = 0; i < length; i++)
            {
                dest[i, index] = src[i];
            }
        }

    }
}
